import numpy as np
import matplotlib.pyplot as plt

arr_rowsum = []

def AMPD(data):
    """
    实现AMPD（Automatic Multi-Scale Peak Detection）算法
    算法过程：首先构造若干个不同尺寸（1-N/2）的窗口，使用不同的窗口在样本数据从左往右滑动，比较[i - k] [i] [i + k]三个值的关系。理论上，峰值的位置，公式[i - k] <= [i] && [i] >= [i + k] 总成立，其它位置部分成立。因此，对每个尺寸窗口成立次数计数，计数最大的那个就是窗口最合适的，也就是峰值所在的位置（计数次数值）。
    这个算法，原理上认为只有一个峰值，但是实际可以计算出任意多个峰值。
    算法缺陷：算法假定样本数据只包含半个周期数据（只有一个峰值，且在数据中间位置），对于挨近数据两头的峰值反而检测不出来
    :param data: 1-D numpy.ndarray 
    :return: 波峰所在索引值的列表
    """
    p_data = np.zeros_like(data, dtype=np.int32)
    count = data.shape[0]

    for k in range(1, count // 2):
        row_sum = 0
        for i in range(k, count - k):
            if data[i] > data[i - k] and data[i] > data[i + k]:
                row_sum -= 1
        print(k, row_sum)
        arr_rowsum.append(row_sum)

    print(arr_rowsum)
    min_index = np.argmin(arr_rowsum)
    max_window_length = min_index
    print(max_window_length)
    for k in range(1, max_window_length + 1):
        print(k, count-k)
        for i in range(k, count - k):
            if data[i] > data[i - k] and data[i] > data[i + k]:
                p_data[i] += 1
    print(p_data)
    print(np.where(p_data == max_window_length))
    return np.where(p_data == max_window_length)[0]

def AMPD2(data):
    length = data.shape[0]
    L = length // 2
    histogram = [0]*length
    
    M = []
    for k in range(0, L):
        M.append([1]*length)
    print(M)

    for k in range(1, L):
        for i in range(k, length-k):
            if (data[i] <= data[i-k]) or (data[i] <= data[i+k]):
                M[k][i] = 0
    print(M)
    
    for i in range(0, length):
        for k in range(0, L):
            histogram[i] += M[k][i]
    print(histogram)

    h_sum = 0
    sq_sum = 0
    count = 0
    for i in range(0, length):
        if histogram[i] > 0:
            h_sum += histogram[i]
            sq_sum += histogram[i] * histogram[i]
            count += 1

    mean = (h_sum / count) if count > 0 else 0
    std_dev = np.sqrt(sq_sum / count - mean * mean) if count > 0 else 0

    print(mean, std_dev)
    threadhold = int(mean + 1*std_dev)
    print(threadhold)

    peaks = []
    peak_count = [0]*length
    
    for i in range(1, length-1):
        if (histogram[i] > threadhold) and (data[i] > data[i - 1]) and (data[i] > data[i + 1]):
            peaks.append(i)
    
    print(peaks, len(peaks))
    return peaks

def AMPD3(data):
    """
    改进 AMPD 算法
    假定：波动周期在 500-7200 之间
    算法过程：
    1. 首先计算样本数据的平均值；
    2. 然后使用平均值查找过均值点（或减去均值）；
    3. 通过过零点（过均值点）将样本数据划分成段；
    4. 针对每一段数据，构造对应大小窗口，使用不同的窗口在样本数据从左往右滑动，比较 [i - k] [i] [i + k] 三个值的关系。理论上，峰值的位置，公式 [i - k] <= [i] && [i] >= [i + k] 总成立，其它位置部分成立。因此，对每个尺寸窗口成立次数计数，计数最大的那个就是窗口最合适的，也就是峰值所在的位置（计数次数值）。
    这个算法，原理上认为只有一个峰值，但是实际可以计算出任意多个峰值。
    算法缺陷：算法假定样本数据只包含半个周期数据（只有一个峰值，且在数据中间位置），对于挨近数据两头的峰值反而检测不出来
    :param data: 1-D numpy.ndarray 
    :return: 波峰所在索引值的列表
    """

    positive = False
    hp_min = 5
    L = data.shape[0]
    avg = np.mean(data)

    if data[0] > avg:
        positive = True
    
    segs = []
    sta = 0
    for i in range(0, L, 3):
        if positive == True:
            if data[i] < avg:
                positive = False
                if (i-sta) > hp_min:
                    segs.append((sta, i))
                    sta = i
        else:
            if data[i] > avg:
                positive = True
                if (i-sta) > hp_min:
                    segs.append((sta, i))
                    sta = i

    if sta != i:
        segs.append((sta, i))

    peaks = []
    for sta, i in segs:
        p = data[sta:i]
        print(p, np.max(p), np.min(p), np.argmax(p), np.argmin(p))
        max_idx = np.argmax(p)
        min_idx = np.argmin(p)
        if max_idx != 0 and max_idx != (len(p)-1):
            peaks.append(max_idx+sta)
        elif min_idx != 0 and min_idx != (len(p)-1):
            peaks.append(min_idx+sta)

    print(peaks)

    return peaks

class DataNode():
    def __init__(self, ts: int, value: float):
        self.ts = ts
        self.value = value

class PeakDetector():
    def __init__(self, period):
        pass
        self._len = 0                   # 当前周期累计计数
        self._sum = 0.0                 # 当前周期累计值
        self._avg = 0.0                 # 上一周期平均值
        self.positive = 0               # 1： 当前检测正周期；-1： 当前检测负周期
        self.hpWin_min = period // 4    # ms 最小检测窗口尺寸（4分之一波）

        self.posPeak = DataNode(0, 0.0)
        self.negPeak = DataNode(0, 0.0)
        self.staPoint = DataNode(0, 0.0)
        self.lastPoint = DataNode(0, 0.0)

    def prepareDetector(self, samp):
        pass
        for s in samp:
            self._sum += s.value
            if self.posPeak.value > s.value:
                self.posPeak = s
            if self.negPeak.value < s.value:
                self.negPeak = s

        self._len += len(samp)
        self._avg = self._sum / self._len

        if samp[-1].value > self._avg:
            self.positive = 1
        else:
            self.positive = -1

        self.staPoint = samp[-1]

    def cleanDetector(self):
        peaks = []

        if (self.lastPoint.ts - self.staPoint.ts) < self.hpWin_min:
            return []

        if self.positive == 1:
            if abs((self.posPeak.value - self._avg) / self._avg) < 0.001:
                return []
            peaks.append(self.posPeak)
        else:
            if abs((self._avg - self.negPeak.value) / self._avg) < 0.001:
                return []
            peaks.append(self.negPeak)

        return peaks

    def detectPeaks(self, samp):
        # 波形限制，最大周期长度内，未找到明显波动波形
        if self._len > 7200*10:
            self._avg = self._sum / self._len
            self._sum = 0.0
            self._len = 0

        peaks = []
        for s in samp:
            self.lastPoint = s
            if self.positive == 1:
                # 查找当前正峰值
                if self.posPeak.value < s.value:
                    self.posPeak = s
                while True:
                    if s.value >= self._avg:
                        # 正周期继续
                        break
                    if abs((self.posPeak.value - self._avg) / self._avg) < 0.001:
                        # 波形幅值限制，幅值不明显，不认为是有效波动
                        break
                    
                    if (self.lastPoint.ts - self.staPoint.ts) > self.hpWin_min:
                        peaks.append(self.posPeak)
                        self.staPoint = s
                    # 切换波峰符号
                    self.positive = -1
                    self.negPeak = s
                    break
            else:
                # 查找当前负峰值
                if self.negPeak.value > s.value:
                    self.negPeak = s
                while True:
                    if s.value <= self._avg:
                        # 负周期继续
                        break
                    if abs((self._avg - self.negPeak.value) / self._avg) < 0.001:
                        # 波形幅值限制，幅值不明显，不认为是有效波动
                        break

                    if (self.lastPoint.ts - self.staPoint.ts) > self.hpWin_min:
                        peaks.append(self.negPeak)
                        self.staPoint = s
                    # 切换波峰符号
                    self.positive = 1
                    self.posPeak = s
                    # 计算上一个周期均值
                    self._avg = self._sum / self._len
                    # 准备下一个周期累计
                    self._sum = 0.0
                    self._len = 0
                    break

            # 累计当前周期
            self._sum += s.value
            self._len += 1

        pass
        return peaks

def sim_data(cycle, N):
    x = np.linspace(0, 2*np.pi * cycle, N)
    y = 5 * np.sin(x)
    noise = np.random.normal(0, 0.1, N)
    y = y + noise

    return y

def vis():
    period = 50
    cycle = 2
    N = period*cycle

    y = sim_data(cycle, N)
    x = range(len(y))
    # print(y)
    plt.plot(range(len(y)), y)
    plt.scatter(range(len(y)), y, color="cyan", alpha=0.6)
    # px = AMPD(y)
    # plt.scatter(px, y[px], color="red")

    # px = AMPD2(y)
    # plt.scatter(px, y[px], color="red")

    # px = AMPD3(y)
    # plt.scatter(px, y[px], color="red")
    
    peaks = []
    sampData = []
    for t, v in zip(x, y):
        node = DataNode(t*10, v)
        sampData.append(node)

    peakDtor = PeakDetector(period*10)
    peakDtor.prepareDetector(sampData)

    sampData = []
    for t, v in zip(x, y):
        node = DataNode(t*10+10*N, v)
        sampData.append(node)
    peaks += peakDtor.detectPeaks(sampData)
    peaks += peakDtor.cleanDetector()

    px = []
    for p in peaks:
        print("p: {:d}; {:.3f}".format(p.ts, p.value))
        px.append((p.ts - 10*N)//10)

    plt.scatter(px, y[px], color="red")
    plt.show()

vis()
